# Dickstein, Ho, and Mark (2023)

get_combined_llh <- function(
  theta, 
  data,
  covar_names,
  covar_lens,
  cost_covar_names,
  cost_covar_lens,
  gradient=TRUE,
  scores=FALSE,
  fixed_cost_censor=NULL){
  
  if (any(is.nan(theta))) {
    print("There are NaNs in the parameter vector")
    return(
      list(
        "objective" = Inf,
        "gradient"  = rep(Inf, length(covar_names))))
  } else {
    
    # The choice bit of the LLH
    choice_theta <- theta
    choice_data <- parametrize_choice_data(choice_theta, data, covar_names, covar_lens)
    choice_llh <- as.numeric(
      choice_data %>%
      mutate(log_s_ij = log(s_ij)) %>%
      mutate(llh_ij = choice * log_s_ij)   %>%
      summarise(sum(llh_ij)))
    choice_score <- choice_data %>% filter(choice == 1) %>% select(starts_with("score_"))
    choice_grad <- colSums(choice_score)    

    # The cost bit of the LLH
    cost_theta <- theta[1:(cost_covar_lens[1]+cost_covar_lens[2])]
    cost_data <- parametrize_cost_data(
      cost_theta,
      data,
      cost_covar_names,
      cost_covar_lens, 
      fixed_cost_censor=fixed_cost_censor)  
    cost_score <- cost_data %>% select(starts_with("cost_score_"))
    cost_llh <- as.numeric(
      cost_data %>% 
      filter(x_av > 0) %>%
      mutate(llh_ij = log_f_c_ij) %>%
      summarise(sum(llh_ij)))
    cost_grad <- colSums(cost_data %>% filter(x_av > 0) %>% select(starts_with("cost_score_")))

    # Extend the cost gradient to have the same number of variables as the choice gradient
    full_cost_grad <- c(cost_grad, rep(0, length(covar_names) - length(cost_covar_names)))

    grad <- choice_grad + full_cost_grad
    llh <- choice_llh + cost_llh
    
    if (is.nan(llh)) {
      llh=-Inf
    }
    if (scores==FALSE) {
      return(
        list(
          "objective" = -llh,
          "gradient"  = -grad))
    } else {
      return(
        list(
          "objective" = -llh,
          "gradient"  = -grad,
          "choice_score" = choice_score,
          "cost_score" = cost_score))
    }
  }
  return(-llh)
}

# Get likelihood
get_choice_llh <- function(
  theta, 
  data,
  covar_names,
  covar_lens,
  gradient=TRUE) {

  if (all(is.nan(theta))) {
    llh <- NaN 
    grad <- theta
    return(
      list(
        "objective" = -llh,
        "gradient"  = -grad))
  } else {
    data <- parametrize_choice_data(theta, data, covar_names, covar_lens)
    llh <- as.numeric(
      data %>%
      mutate(log_s_ij = log(s_ij)) %>%
      mutate(llh_ij = choice * log_s_ij)   %>%
      summarise(sum(llh_ij)))
    grad <- colSums(data %>% select(starts_with("score_")))
    if (is.nan(llh)) {
      llh <- -Inf
    }

    if (gradient==FALSE) {
      return(-llh)
    }
    else {
      return(
        list(
          "objective" = -llh,
          "gradient"  = -grad))
    }
    }
}

get_cost_llh <- function(
  cost_theta, 
  data,
  cost_covar_names,
  cost_covar_lens,
  gradient=TRUE,
  fixed_cost_censor=NULL) {

  cost_data <- parametrize_cost_data(
    cost_theta,
    data,
    cost_covar_names, 
    cost_covar_lens,
    fixed_cost_censor=fixed_cost_censor) %>%
  filter(choice == 1, x_av > 0) # Only insured people contribute to the cost LLH
  
  llh <- as.numeric(
    cost_data %>%
    mutate(llh_ij = log_f_c_ij)   %>%
    summarise(sum(llh_ij)))
  grad <- colSums(cost_data %>% select(starts_with("cost_score_")))
  if (is.nan(llh)) {
    llh <- -Inf
  }

  if (gradient==FALSE) {
    return(-llh)
  }
  else {
    return(
      list(
        "objective" = -llh,
        "gradient"  = -grad))
  }
}
